import torch
from torch import nn
from transformers import AutoTokenizer, AutoConfig

from src.sdpa2attention import MultiHeadAttention


# 定义前馈网络（FeedForward）类，继承自 PyTorch 的 nn.Module
class FeedForward(nn.Module):
    def __init__(self, config):
        """
        初始化前馈网络模块。

        参数：
        - config: 配置对象，包含以下属性：
            - hidden_size: 输入和输出的特征维度（embedding dimension）。
            - intermediate_size: 前馈网络中间层的特征维度。
            - hidden_dropout_prob: Dropout 的概率，用于防止过拟合。
        """
        super().__init__()  # 调用父类的初始化方法

        # 定义第一个线性层，将输入从 hidden_size 映射到 intermediate_size
        self.linear_1 = nn.Linear(config.hidden_size, config.intermediate_size)

        # 定义第二个线性层，将中间层的输出映射回 hidden_size
        self.linear_2 = nn.Linear(config.intermediate_size, config.hidden_size)

        # 定义激活函数，使用 GELU（Gaussian Error Linear Unit）
        # GELU 是 Transformer 中常用的激活函数，效果优于 ReLU
        self.gelu = nn.GELU()

        # 定义 Dropout 层，用于在训练时随机丢弃部分神经元，防止过拟合
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, x):
        """
        前向传播函数，计算前馈网络的输出。

        参数：
        - x: 输入张量，形状为 [batch_size, seq_len, hidden_size]

        返回：
        - x: 输出张量，形状为 [batch_size, seq_len, hidden_size]
        """
        # 第一个线性变换：将输入从 hidden_size 映射到 intermediate_size
        x = self.linear_1(x)

        # 应用 GELU 激活函数，增加非线性
        x = self.gelu(x)

        # 第二个线性变换：将中间层的输出映射回 hidden_size
        x = self.linear_2(x)

        # 应用 Dropout，防止过拟合
        x = self.dropout(x)

        # 返回最终的输出
        return x

class Embeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.token_embeddings = nn.Embedding(config.vocab_size,
                                             config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings,
                                                config.hidden_size)
        self.layer_norm = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout()

    def forward(self, input_ids):
        # Create position IDs for input sequence
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long).unsqueeze(0)
        # Create token and position embeddings
        token_embeddings = self.token_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        # Combine token and position embeddings
        embeddings = token_embeddings + position_embeddings
        embeddings = self.layer_norm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings




# 定义 Transformer 编码器层（TransformerEncoderLayer），继承自 PyTorch 的 nn.Module
class TransformerEncoderLayer(nn.Module):
    def __init__(self, config):
        """
        初始化 Transformer 编码器层。

        参数：
        - config: 配置对象，包含以下属性：
            - hidden_size: 输入和输出的特征维度（embedding dimension）。
            - num_heads: 多头注意力的头数。
            - intermediate_size: 前馈网络中间层的特征维度。
            - hidden_dropout_prob: Dropout 的概率，用于防止过拟合。
        """
        super().__init__()  # 调用父类的初始化方法

        # 定义第一个 LayerNorm，用于对输入进行归一化
        self.layer_norm_1 = nn.LayerNorm(config.hidden_size)

        # 定义第二个 LayerNorm，用于对前馈网络的输入进行归一化
        self.layer_norm_2 = nn.LayerNorm(config.hidden_size)

        # 定义多头注意力模块
        self.attention = MultiHeadAttention(config)

        # 定义前馈网络模块
        self.feed_forward = FeedForward(config)  #- config intermediate_size: 前馈网络中间层的特征维度。  - hidden_dropout_prob: Dropout 的概率，用于防止过拟合。

    def forward(self, x, mask=None):
        """
        前向传播函数，计算 Transformer 编码器层的输出。

        参数：
        - x: 输入张量，形状为 [batch_size, seq_len, hidden_size]。
        - mask: 注意力掩码（可选），形状为 [batch_size, seq_len, seq_len]，用于屏蔽某些位置的注意力。

        返回：
        - x: 输出张量，形状为 [batch_size, seq_len, hidden_size]。
        """
        # 第一步：对输入进行 LayerNorm 归一化
        hidden_state = self.layer_norm_1(x)

        # 第二步：多头注意力机制
        # 输入的 query、key 和 value 都是归一化后的 hidden_state
        # 使用残差连接（skip connection）：将注意力输出与原始输入相加
        x = x + self.attention(hidden_state, hidden_state, hidden_state, mask=mask)

        # 第三步：前馈网络
        # 对注意力输出进行第二次 LayerNorm 归一化
        # 使用残差连接：将前馈网络的输出与注意力输出相加
        x = x + self.feed_forward(self.layer_norm_2(x))

        # 返回最终的输出
        return x
class TransformerEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.embeddings = Embeddings(config)
        self.layers = nn.ModuleList([TransformerEncoderLayer(config)
                                     for _ in range(config.num_hidden_layers)])

    def forward(self, x, mask=None):
        x = self.embeddings(x)
        for layer in self.layers:
            x = layer(x, mask=mask)
        return x


if __name__ == '__main__':
    model_ckpt = "bert-base-uncased"

    # 加载与指定模型对应的分词器
    tokenizer = AutoTokenizer.from_pretrained(model_ckpt)

    # 定义输入文本
    text = "time flies like an arrow"
    inputs = tokenizer(text, return_tensors="pt", add_special_tokens=False)
    config = AutoConfig.from_pretrained(model_ckpt)
    token_emb = nn.Embedding(config.vocab_size, config.hidden_size)  # 初始化embedding class
    inputs_embeds = token_emb(inputs.input_ids) #使用embedding


    #前馈网络
    multihead_attn = MultiHeadAttention(config)
    query = key = value = inputs_embeds
    attn_output = multihead_attn(query, key, value)
    feed_forward = FeedForward(config)
    ff_outputs = feed_forward(attn_output)
    print(ff_outputs.size())

    encoder_layer = TransformerEncoderLayer(config)
    print(f"inputs_embeds.shape“{inputs_embeds.shape}")
    print(f"encoder_layer(inputs_embeds).shape,{encoder_layer(inputs_embeds).shape}")

    embedding_layer = Embeddings(config)
    print(f"embedding_layer(inputs.input_ids).shape, {embedding_layer(inputs.input_ids).shape}")

    encoder = TransformerEncoder(config)
    print(f"encoder(inputs.input_ids).shape, {encoder(inputs.input_ids).shape}")